package org.example.common.threadpool;

import com.alibaba.ttl.threadpool.TtlExecutors;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.boot.autoconfigure.task.TaskExecutionProperties;
import org.springframework.boot.task.TaskExecutorBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.lang.reflect.Method;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * 配置线程池，使用不同线程执行任务，提升效率
 *
 * @author dingcan
 * @since 2019-12-16
 */
@Slf4j
@EnableAsync
@Configuration(proxyBeanMethods = false)
@RequiredArgsConstructor
public class ThreadPoolConfig implements AsyncConfigurer {
    public static final String ASYNC_EXECUTOR_NAME = "taskExecutor";
    private final TaskExecutionProperties taskExecutionProperties;
    
    /**
     * 异步线程池
     */
    @Bean(name = ASYNC_EXECUTOR_NAME)
    public ExecutorService threadPoolTaskExecutor() {
        return TtlExecutors.getTtlExecutorService(createThreadPool(taskExecutionProperties, new ThreadPoolExecutor.CallerRunsPolicy()));
    }
    
    public static ExecutorService createThreadPool(TaskExecutionProperties taskExecutionProperties, RejectedExecutionHandler rejectedExecutionHandler) {
        // 创建 TaskExecutorBuilder 对象
        TaskExecutorBuilder builder = createTskExecutorBuilder(taskExecutionProperties);
        // 创建 ThreadPoolTaskExecutor 对象
        ThreadPoolTaskExecutor threadPoolTaskExecutor = builder.build();
        // 配置拒绝策略
        // AbortPolicy：默认的拒绝策略，会throws RejectedExecutionException。
        // CallerRunsPolicy：提交任务的线程自己去执行该任务。如果执行程序已关闭，则丢弃。
        // DiscardOldestPolicy：丢弃最老的任务，其实就是把最早进入工作队列的任务丢弃，然后把新任务加入到工作队列。
        // DiscardPolicy：直接丢弃任务，没有任何异常抛出。
        threadPoolTaskExecutor.setRejectedExecutionHandler(rejectedExecutionHandler);
        threadPoolTaskExecutor.setDaemon(true);
        // 线程池初始化
        threadPoolTaskExecutor.initialize();
        log.info("线程池初始化成功:{}", taskExecutionProperties.getThreadNamePrefix());
        return threadPoolTaskExecutor.getThreadPoolExecutor();
    }
    
    
    /**
     * 异步线程池配置
     *
     * @param properties 配置参数
     * @return task executor builder
     */
    public static TaskExecutorBuilder createTskExecutorBuilder(TaskExecutionProperties properties) {
        // Pool 属性
        TaskExecutionProperties.Pool pool = properties.getPool();
        TaskExecutionProperties.Shutdown shutdown = properties.getShutdown();
        return new TaskExecutorBuilder()
                // 缓冲队列任务个数
                .queueCapacity(pool.getQueueCapacity())
                // 核心线程池，默认值1
                .corePoolSize(pool.getCoreSize())
                // 线程池最大值，当核心线程池被用完且缓冲队列已满时，才会申请
                .maxPoolSize(pool.getMaxSize())
                // 核心线程是否允许超时
                .allowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout())
                // 超出核心线程池之外的空闲线程最大存活时间
                .keepAlive(pool.getKeepAlive())
                // 线程名前缀
                .threadNamePrefix(properties.getThreadNamePrefix())
                // Shutdown 属性
                .awaitTermination(shutdown.isAwaitTermination())
                .awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod())
                // 其它基本属性
                .threadNamePrefix(properties.getThreadNamePrefix())
                // 异步线程复制上下文信息
                .taskDecorator(new ContextCopyingDecorator());
    }
    
    /**
     * 异步任务异常处理
     * 只处理没有返回值的异步任务
     * 有返回值的异步任务，由客户端自行处理
     */
    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return new AsyncExceptionHandler();
    }
    
    private static class AsyncExceptionHandler implements AsyncUncaughtExceptionHandler {
        
        /**
         * 处理未捕获异常
         *
         * @param ex     异步任务抛出的异常对象
         * @param method 异步任务处理的方法
         * @param params 异常处理需要传递的参数
         */
        @Override
        public void handleUncaughtException(Throwable ex, Method method, Object... params) {
            
            log.warn("线程池报错,  AsyncError: {}, Method: {},  Param: {},  StackTrace:",
                    ex.getMessage(),
                    method.getName(),
                    params,
                    ex);
        }
    }
}
